import joblib  # 用于加载和保存文件
import tensorflow as tf 
import argparse  # 用于解析命令行参数
import os.path as osp 
from learning_to_adapt.samplers.utils import rollout 
import json  

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    # pkl 和 json 文件所在的目录
    parser.add_argument("param", type=str, help='Directory with the pkl and json file')
    # -l，最大路径长度
    parser.add_argument('--max_path_length', '-l', type=int, default=1000,
                        help='Max length of path')
    # num_rollouts 缩写为-n 最大轮数
    parser.add_argument('--num_rollouts', '-n', type=int, default=10,
                        help='Max length of rollout')
    # 加速比？
    parser.add_argument('--speedup', type=float, default=1,
                        help='Speedup')
    # 输出视频的目录
    parser.add_argument('--video_filename', type=str,
                        help='path to the out video file')
    # 是否进行更多模拟？
    parser.add_argument('--prompt', type=bool, default=False,
                        help='Whether or not to prompt for more sim')
    # 当环境完成时是否停止动画
    parser.add_argument('--ignore_done', action='store_true',
                        help='Whether stop animation when environment done or continue anyway')
    # 解析命令行参数
    args = parser.parse_args()

    with tf.Session() as sess:
        # 构建pkl文件路径
        pkl_path = osp.join(args.param, 'params.pkl')
        # 构建json文件路径
        json_path = osp.join(args.param, 'params.json')
        print("Testing policy %s" % pkl_path)
        # 加载并解析json文件内容到json_params变量
        json_params = json.load(open(json_path, 'r'))
        # 使用joblib加载pkl文件内容到data变量
        data = joblib.load(pkl_path)
        # 从data中获取policy和env
        policy = data['policy']
        env = data['env']
        # 进行rollouts(模拟环境的多次迭代)操作，执行args.num_rollouts次
        for _ in range(args.num_rollouts):
            path = rollout(env, policy, max_path_length=args.max_path_length,
                           animated=True, ignore_done=args.ignore_done,
                           adapt_batch_size=json_params.get('adapt_batch_size', None))
            # 打印路径上的总奖励和
            # print(sum(path['rewards']))
